import numpy as np
import random


class TSPEnv(object):

    def __init__(self, _points_num, _distance_matrix):
        # 一共多少个点
        self.points_num = _points_num
        # 动作就有几种
        self.action_space = self.points_num
        self.action_select_list = [i for i in range(0, self.action_space)]
        # 那么状态就有几种
        self.observation_space = self.points_num
        # 记录走过的路,其实历史路劲就是observation
        self.stops = []
        # 距离矩阵，类型numpy,以点的顺序作为矩阵中距离的顺序，比如
        # a=np.assarry([[1, 2, 3],
        # [4, 1, 6],
        # [7, 8, 1]])
        # 这里第0个点到第1个点的距离就是a[0, 1],第1个点到第0个点就是a[1,0]，可以不一样（来去“距离”车费不一样）
        self.distance_matrix = _distance_matrix

    def render(self):
        # 打印走过的路，如果想要画图就在这里
        out_str = ''
        for i in self.stops:
            out_str = out_str + str(i) + '-'
        print (out_str)
        return out_str,self.stops

    def reset(self):
        # 重置路劲数据
        self.stops = []
        # 随机生成第一个初始节点,行动数组里随便找一个点
        randi = random.randint(0, self.points_num - 1)
        first_stop = self.action_select_list[randi]
        self.stops.append(first_stop)
        return first_stop

    def step(self, next_step):
        # 历史数组中的上一步即是当前状态
        this_state = self.stops[-1]
        new_state = next_step
        # 将新的路添加进历史路劲
        self.stops.append(next_step)

        # 奖励得分，目前是按距离
        reward = self.getReward(this_state, new_state)

        done = False
        # 终止条件，如果历史路劲数组的长度和输入的地点的长度一样，则表示已完成（TSP问题不走回头路），
        # 也可以加入走回头路就是done的判断段
        if len(self.stops) == self.points_num:
            done = True

        return new_state, reward, done

    def getReward(self, this_state, new_state):
        reward = self.distance_matrix[this_state, new_state]
        return reward


if __name__ == '__main__':
    distance_matrix = np.asarray(
        [
            [1, 2, 3, 4, 5],
            [2, 1, 3, 4, 5],
            [3, 3, 1, 4, 5],
            [4, 4, 4, 1, 5],
            [5, 5, 5, 5, 1]
        ]
    )

    tspEnv = TSPEnv(_points_num=5, _distance_matrix=distance_matrix)
    first_stop = tspEnv.reset()
    print(first_stop)

    print(np.random.rand())
    pass
